import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

def visualize_knn_boundary(knn_model, X_train, y_train, X_test=None, y_test=None, feature_indices=(0, 1)):
    """
    可视化 KNN 分类器的决策边界

    参数:
        knn_model: 已训练的 KNN 模型对象（支持 `predict` 方法）
        X_train (ndarray): 训练数据 (N, M)，其中 N 是样本数量，M 是特征数量
        y_train (ndarray): 训练数据的标签 (N,)
        X_test (ndarray, optional): 测试数据 (N_test, M)，默认为 None
        y_test (ndarray, optional): 测试数据的标签 (N_test,)，默认为 None
        feature_indices (tuple, optional): 用于可视化的两个特征的索引 (i, j)，默认为 (0, 1)

    输出:
        决策边界图
    """
    # 选择用于可视化的两个特征
    X_train_2d = X_train[:, feature_indices]
    if X_test is not None:
        X_test_2d = X_test[:, feature_indices]

    # 定义网格范围
    x_min, x_max = X_train_2d[:, 0].min() - 1, X_train_2d[:, 0].max() + 1
    y_min, y_max = X_train_2d[:, 1].min() - 1, X_train_2d[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.01), np.arange(y_min, y_max, 0.01))

    # 预测网格上的分类结果
    Z = knn_model.predict(np.c_[xx.ravel(), yy.ravel()])
    Z = Z.reshape(xx.shape)

    # 定义颜色映射
    cmap_light = ListedColormap(['#FFAAAA', '#AAAAFF', '#AAFFAA'])
    cmap_bold = ListedColormap(['#FF0000', '#0000FF', '#00FF00'])

    # 绘制决策边界
    plt.figure(figsize=(10, 6))
    plt.contourf(xx, yy, Z, alpha=0.8, cmap=cmap_light)

    # 绘制训练集和测试集样本点
    plt.scatter(X_train_2d[:, 0], X_train_2d[:, 1], c=y_train, cmap=cmap_bold, edgecolor='k', s=20, label="Train")
    if X_test is not None:
        plt.scatter(X_test_2d[:, 0], X_test_2d[:, 1], c=y_test, cmap=cmap_bold, marker='x', s=50, label="Test")

    # 添加图例和标题
    plt.title("KNN Decision Boundary")
    plt.xlabel(f"Feature {feature_indices[0]}")
    plt.ylabel(f"Feature {feature_indices[1]}")
    plt.legend()
    plt.grid(alpha=0.7)
    plt.show()